
#include <iostream>
#include <memory>

#include "helper/helper.h"
#include "data/MemoryDataStore.h"
#include "data/StochasticDataFeeder.h"
#include "algorithm/Perceptron.h"
#include "algorithm/RegularizedPerceptron.h"
#include "algorithm/RDAPerceptron.h"
#include "algorithm/CommonLoss.h"
#include "algorithm/CommonRegularizer.h"
#include "algorithm/StopWatch.h"

po::variables_map vm;
po::options_description desc("Allowed options");

// JOBS
void perceptron_learn();

int main(int argc, char** argv)
{
  //std::cerr << "Online learning" << std::endl;

  desc.add_options()
    ("help", "produce help message")
    // DATA
    ("data", po::value<std::string>(), "data file")
    ("test", po::value<std::string>(), "test file")
    ("hint", po::value<int>()->default_value(0), "data dimension hint")
    // MODEL
    ("model", po::value<std::string>(), "output model file")
    // METHODS
    ("perceptron", "classical perceptron learning\n--rate")
    ("regularized", "regularized perceptron learning\n--rate --loss --reg --C --iter")
    ("rda", "regularized dual averaging perceptron learning\n--rate --loss --reg --C --iter")
    // COMMON PARAMETERS
    ("iter", po::value<int>(), "maximum iterations")
    ("rate", po::value<double>()->default_value(1.0), "learning rate")
    ("loss", po::value<std::string>(), "loss function [hinge,log,exp]")
    ("reg", po::value<std::string>(), "regularizing function [l1,l2]")
    ("C", po::value<double>()->default_value(1.0), "weight on new sample")
    ;

  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm); 

  if (vm.count("help")) {
    std::cerr << desc << std::endl;
  }
  else if (vm.count("perceptron") || vm.count("regularized") || vm.count("rda")) {
    perceptron_learn();
  }

  return 0;
}

// Create learner based on input parameters (command line)
std::shared_ptr<StochasticAlgorithm> CreateLearner() 
{
  PO_VARIABLE(rate, vm, "rate", double);
  if (vm.count("perceptron")) 
  {
    //return std::shared_ptr<StochasticAlgorithm>(new Perceptron(rate));
    return std::make_shared<Perceptron>(rate);
  }
  else if (vm.count("regularized")) 
  {
    PO_REQUIRE(vm, "loss");
    PO_REQUIRE(vm, "reg");

    PO_VARIABLE(C, vm, "C", double);
    PO_VARIABLE(loss, vm, "loss", std::string);
    PO_VARIABLE(reg, vm, "reg", std::string);

    if (loss == "hinge" && reg == "l2")
      return std::make_shared<RegularizedPerceptron<HingeLoss,L2Regularizer> >(rate, C);
    //return std::shared_ptr<StochasticAlgorithm>(new RegularizedPerceptron<HingeLoss,L2Regularizer>(rate, C));
    else if (loss == "hinge" && reg == "l1")
      return std::make_shared<RegularizedPerceptron<HingeLoss,L1Regularizer> >(rate, C);
    else {
      std::cerr << "Unknown loss and regularizer combination" << std::endl;
      exit(1);
    }
  }
  else if (vm.count("rda")) 
  {
    PO_REQUIRE(vm, "loss");
    PO_REQUIRE(vm, "reg");

    PO_VARIABLE(C, vm, "C", double);
    PO_VARIABLE(loss, vm, "loss", std::string);
    PO_VARIABLE(reg, vm, "reg", std::string);

    if (loss == "hinge" && reg == "l2")
      return std::make_shared<RDAPerceptron<HingeLoss,L2Regularizer> >(rate, C);
    else {
      std::cerr << "Unknown loss and regularizer combination" << std::endl;
      exit(1);
    }
  }
  else {
    std::cerr << "Unknown learner type should be --perceptron" << std::endl;
    exit(1);
  }
}

// Perceptron learning and testing
void perceptron_learn()
{
  StopWatch watch;

  PO_REQUIRE(vm, "data");
  PO_REQUIRE(vm, "iter");

  PO_VARIABLE(dataFile, vm, "data", std::string);
  PO_VARIABLE_DEFAULT(testFile, vm, "test", std::string, dataFile);
  PO_VARIABLE(maxIter, vm, "iter", int);
  PO_VARIABLE(hint, vm, "hint", int);

  // Read a data file, currently in LIBSVM format
  MemoryDataStore dataStore(dataFile, DataFormat::LIBSVM, hint);
  //std::cerr << "data read d=" << watch.GetDuration() << std::endl;
  StochasticDataFeeder<MemoryDataStore> dataFeeder(dataStore);

  std::shared_ptr<StochasticAlgorithm> pAlg = CreateLearner();

  double errorCount = 0;

  MemoryDataStore testDataStore(testFile, DataFormat::LIBSVM, hint);
  //std::cerr << "test data read d=" << watch.GetDuration() << std::endl;
  for (int iter = 0; iter < maxIter; iter++) {
    const Example& example = dataFeeder.GetNextExample();
    double predict_label = pAlg->Learn(example,iter);
    if (predict_label != example.GetLabel()) errorCount++;

    if (iter % 100 == 0 || iter == maxIter - 1) {
      double testErrorCount = 0;
      for (auto& example : testDataStore.GetExamples()) {
        double predict_label = pAlg->Predict(example.GetInput());
        if (predict_label != example.GetLabel()) testErrorCount++;
      }
      std::cerr << iter << "\t " << (errorCount/(iter+1)*100) << "\t" << 
        testErrorCount/testDataStore.GetExampleCount()*100 << 
        "\t" << watch.GetTotalDuration() << std::endl;
    }
  }

  std::cout << errorCount << std::endl;
  std::cerr << "all done d=" << watch.GetTotalDuration() << std::endl;
  if (vm.count("model")) {
    PO_VARIABLE(modelFile, vm, "model", std::string);
    pAlg->Save(modelFile);
  }
}
